# -*- coding: utf-8 -*-
"""
HaaS506 RTU OTA升级模块
提供远程固件和脚本升级功能
"""

import utime
import ujson as json
import uos
import usocket as socket
import http
import modem
# System functions
import rtu_config
import rtu_mqtt
import rtu_system

# OTA状态
_ota_in_progress = False
_last_check_time = 0
_check_interval = 3600  # 1小时检查一次
_update_available = False
_update_info = {}
_module_status = "STOPPED"

def init():
    """初始化OTA模块"""
    global _module_status
    _module_status = "RUNNING"
    print("[OTA] OTA module initialized")
    return True

def check_update():
    """检查服务器是否有更新"""
    global _last_check_time, _update_available, _update_info
    
    if _module_status != "RUNNING":
        print("[OTA] OTA module not running")
        return False
    
    try:
        # 从config获取OTA服务器配置
        ota_server = rtu_config.get_config("servers.ota_server.primary_url", "http://dtu.coldwang.com/api/ota/check")
        timeout = rtu_config.get_config("servers.ota_server.timeout", 30)
        
        # 解析服务器URL
        if ota_server.startswith("http://"):
            ota_server = ota_server[7:]
        
        parts = ota_server.split("/", 1)
        host = parts[0]
        path = "/" + (parts[1] if len(parts) > 1 else "ota")
        
        print("[OTA] Checking for updates from: {}{}".format(host, path))
        
        # 创建HTTP请求
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(timeout)
        
        try:
            addr = socket.getaddrinfo(host, 80)[0][-1]
            sock.connect(addr)
        except Exception as e:
            print("[OTA] Socket connect error:", e)
            sock.close()
            return False
        
        # 构建请求数据
        request_data = {
            "action": "check_update",
            "current_version": rtu_config.get_config("system.version", "2.0.0"),
            "device_sn": rtu_config.get_device_sn(),
            "device_type": rtu_config.get_config("device.hardware", "haas506_rtu")
        }
        
        print("[OTA] Request data: {}".format(json.dumps(request_data)))
        
        request = "POST {} HTTP/1.1\r\nHost: {}\r\nContent-Type: application/json\r\nConnection: close\r\nContent-Length: {}\r\n\r\n{}".format(
            path, host, len(json.dumps(request_data)), json.dumps(request_data)
        )
        
        sock.send(request.encode())
        response = sock.read()
        sock.close()
        
        # 解析响应
        response_str = response.decode("utf-8", "ignore")
        
        if "200 OK" not in response_str:
            print("[OTA] Server returned non-200 response")
            return False
        
        # 提取JSON数据
        json_start = response_str.find("\r\n\r\n")
        if json_start == -1:
            json_start = response_str.find("\n\n") + 2
        else:
            json_start += 4
        
        data = json.loads(response_str[json_start:])
        print("[OTA] Server response: {}".format(data))
        
        # 发布状态到MQTT
        if rtu_config.g_mqtt_connected:
            rtu_mqtt.publish_status(data)
        
        # 处理服务器响应
        _update_available = data.get("update_available", False)
        if _update_available:
            _update_info = {
                "version": data.get("version", ""),
                "url": data.get("url", ""),
                "type": data.get("type", "python"),
                "description": data.get("description", "")
            }
            print("[OTA] Update available: {}".format(_update_info["version"]))
        else:
            print("[OTA] No updates available")
        
        _last_check_time = utime.time()
        return _update_available
        
    except Exception as e:
        print("[OTA] Update check failed: {}".format(e))
        # 发布错误状态
        check_info = {
            "action": "update_check_result",
            "error": str(e),
            "current_version": rtu_config.get_config("system.version", "2.0.0")
        }
        if rtu_config.g_mqtt_connected:
            rtu_mqtt.publish_status(check_info)
        return False

def is_update_needed():
    """是否需要检查更新"""
    global _last_check_time
    
    # 检查更新间隔
    current_time = utime.time()
    if current_time - _last_check_time < _check_interval:
        return False
    
    # 检查网络状态
    if not rtu_config.g_network_ready:
        return False
    
    return True

def download_and_apply_update():
    """下载并应用更新"""
    global _update_available, _update_info, _ota_in_progress
    
    if not _update_available or not _update_info:
        print("[OTA] No update available to apply")
        return False
    
    try:
        print("[OTA] Starting update process...")
        print("[OTA] Update info: {}".format(json.dumps(_update_info)))
        
        _ota_in_progress = True
        
        # 在OTA开始前验证device_sn持久化（如果不存在会自动创建）
        print("[OTA] Verifying device_sn persistence before OTA...")
        if not _verify_device_sn_persistence():
            print("[OTA] WARNING: device_sn persistence verification failed before OTA!")
            # 尝试获取并保存device_sn
            device_sn = rtu_config.get_device_sn()
            if device_sn and str(device_sn) != "unknown" and str(device_sn) != "10000":
                print("[OTA] Attempting to save device_sn before OTA: {}".format(device_sn))
                try:
                    rtu_config.save_device_sn_persistent(device_sn, source="ota")
                    # 再次验证
                    if not _verify_device_sn_persistence():
                        print("[OTA] CRITICAL: Cannot ensure device_sn persistence, aborting OTA!")
                        _ota_in_progress = False
                        return False
                except Exception as e:
                    print("[OTA] Failed to save device_sn: {}".format(e))
                    _ota_in_progress = False
                    return False
            else:
                print("[OTA] CRITICAL: Invalid device_sn ({}), aborting OTA!".format(device_sn))
                _ota_in_progress = False
                return False
        
        # 保存当前版本信息（在任何配置更新之前）
        old_version = rtu_config.get_config("system.version", "2.0.0")
        
        # 创建HTTP客户端 - 使用简化方式
        update_type = _update_info.get("type", "python")
        ota_url = _update_info.get("url", "")
        
        if not ota_url:
            print("[OTA] No update URL provided")
            return False
        
        print("[OTA] Downloading update from: {}".format(ota_url))
        
        # 记录升级信息
        upgrade_info = {
            "action": "ota_upgrade_start",
            "old_version": old_version,
            "new_version": _update_info["version"],
            "type": update_type
        }
        print("[OTA] Upgrade Start: {}".format(json.dumps(upgrade_info)))
        
        if rtu_config.g_mqtt_connected:
            rtu_mqtt.publish_status(upgrade_info)
        
        # 创建HTTP客户端进行实际下载
        htp = http.client()
        
        # 根据更新类型选择不同的升级方法
        if update_type == "python":
            # Python脚本升级
            print("[OTA] Starting Python script OTA...")
            result = htp.fotaPy(ota_url)
        elif update_type == "firmware":
            # 固件升级
            print("[OTA] Starting firmware OTA...")
            result = htp.fotaSys(ota_url)
        else:
            print("[OTA] Unknown update type: {}".format(update_type))
            _ota_in_progress = False
            return False
        
        if result == 0:
            print("[OTA] Update download successful, preparing to reboot...")
            
            # 简化流程：OTA成功后的3步处理
            # 步骤1: 检查并创建 device_sn.txt（如果需要）
            device_sn = None
            try:
                # 先尝试从当前运行的配置获取device_sn
                device_sn = rtu_config.get_device_sn()
                if device_sn and str(device_sn) != "unknown" and str(device_sn) != "10000":
                    # 检查文件是否存在
                    try:
                        with open("/haas/device_sn.txt", "r") as f:
                            saved_sn = f.read().strip()
                        print("[OTA] device_sn.txt already exists: {}".format(saved_sn))
                    except:
                        # 文件不存在，创建它
                        print("[OTA] Creating device_sn.txt with: {}".format(device_sn))
                        try:
                            with open("/haas/device_sn.txt", "w") as f:
                                f.write(str(device_sn))
                            print("[OTA] device_sn.txt created successfully")
                        except Exception as e:
                            print("[OTA] Failed to create device_sn.txt: {}".format(e))
                else:
                    print("[OTA] Invalid device_sn: {}, skip creating file".format(device_sn))
            except Exception as e:
                print("[OTA] Error checking device_sn: {}".format(e))
            
            # 步骤2&3: 配置迁移将在重启后由 rtu_config.load_config() 自动处理
            print("[OTA] Config migration will be handled automatically on next boot")
            
            # 记录升级成功信息
            old_version = rtu_config.get_config("system.version", "2.0.0")
            upgrade_success = {
                "action": "ota_upgrade_success",
                "old_version": old_version,
                "new_version": _update_info["version"],
                "type": update_type,
                "device_sn": device_sn
            }
            print("[OTA] Upgrade Success: {}".format(json.dumps(upgrade_success)))
            
            if rtu_config.g_mqtt_connected:
                rtu_mqtt.publish_status(upgrade_success)
            
            # 延迟后重启
            utime.sleep(3)  # 增加延迟确保文件系统同步
            print("[OTA] Rebooting system...")
            modem.reboot()  # 使用modem.reboot()重启设备
            
            _ota_in_progress = False
            return True
        else:
            print("[OTA] Update download failed with code: {}".format(result))
            
            # 发布下载失败状态
            upgrade_failed = {
                "action": "ota_upgrade_failed",
                "old_version": old_version,
                "new_version": _update_info["version"],
                "type": update_type,
                "error": "Download failed with code: {}".format(result)
            }
            print("[OTA] Upgrade Failed: {}".format(json.dumps(upgrade_failed)))
            
            if rtu_config.g_mqtt_connected:
                rtu_mqtt.publish_status(upgrade_failed)
            
            _ota_in_progress = False
            return False
        
    except Exception as e:
        print("[OTA] Update process failed: {}".format(e))
        upgrade_info = {
            "action": "ota_upgrade_failed",
            "old_version": old_version,
            "new_version": _update_info.get("version", "unknown"),
            "error": str(e)
        }
        if rtu_config.g_mqtt_connected:
            rtu_mqtt.publish_status(upgrade_info)
        
        _ota_in_progress = False
        return False

def _verify_device_sn_persistence():
    """简化的device_sn持久化验证 - 如果device_sn有效但文件不存在则创建"""
    try:
        # 1. 获取当前的device_sn
        current_device_sn = rtu_config.get_device_sn()
        current_sn_str = str(current_device_sn) if current_device_sn else ""
        
        # 2. 检查是否是有效的device_sn
        if not current_sn_str or current_sn_str == "unknown" or current_sn_str == "10000":
            print("[OTA] Invalid device_sn: {}, cannot ensure persistence".format(current_sn_str))
            return False
        
        # 3. 检查持久化文件是否存在
        try:
            with open("/haas/device_sn.txt", "r") as f:
                saved_sn = f.read().strip()
            print("[OTA] device_sn.txt exists with value: {}".format(saved_sn))
            
            # 如果文件存在但值不同，更新文件
            if saved_sn != current_sn_str:
                print("[OTA] Updating device_sn.txt: {} -> {}".format(saved_sn, current_sn_str))
                try:
                    rtu_config.save_device_sn_persistent(current_device_sn, source="ota")
                    print("[OTA] device_sn.txt updated successfully")
                except Exception as e:
                    print("[OTA] Failed to update device_sn.txt: {}".format(e))
                    return False
            
            return True
            
        except:
            # 4. 文件不存在，创建它
            print("[OTA] device_sn.txt not found, creating with value: {}".format(current_sn_str))
            try:
                rtu_config.save_device_sn_persistent(current_device_sn, source="ota")
                
                # 验证创建成功
                with open("/haas/device_sn.txt", "r") as f:
                    saved_sn = f.read().strip()
                print("[OTA] device_sn.txt created successfully with value: {}".format(saved_sn))
                return True
                
            except Exception as e:
                print("[OTA] Failed to create device_sn.txt: {}".format(e))
                return False
        
    except Exception as e:
        print("[OTA] Error in device_sn persistence check: {}".format(e))
        return False


def _deep_merge_configs(base_config, new_config):
    """深度合并配置，保留旧值，只添加新增项"""
    def deep_update(base, update):
        for key, value in update.items():
            if key not in base:
                # 新增项，直接添加
                base[key] = value
                print("[OTA] Added new config key: {}".format(key))
            elif isinstance(value, dict) and isinstance(base.get(key), dict):
                # 递归合并字典
                deep_update(base[key], value)
            # 否则保留旧值（不做任何操作）
    
    merged = base_config.copy()
    deep_update(merged, new_config)
    return merged


def manual_trigger():
    """外部模块调用的OTA强制触发接口"""
    print("[OTA] Manual OTA trigger called...")
    if check_update():
        return download_and_apply_update()
    return False

def force_update_check():
    """强制检查更新"""
    return manual_trigger()

def get_ota_status():
    """获取OTA状态"""
    time_since_check = int(utime.time() - _last_check_time) if _last_check_time > 0 else -1
    return {
        "status": _module_status,
        "current_version": rtu_config.get_config("system.version", "2.0.0"),
        "update_available": _update_available,
        "ota_in_progress": _ota_in_progress,
        "last_check_time": _last_check_time,
        "last_check_ago": "{}s".format(time_since_check) if time_since_check >= 0 else "Never",
        "check_interval": _check_interval
    }

def _notify_ota_status(status, version=None, error=None):
    """通知OTA状态"""
    try:
        data = {
            "action": "ota_status",
            "status": status,
            "timestamp": utime.time()
        }
        
        if version:
            data["version"] = version
        if error:
            data["error"] = error
        
        # 通过MQTT发送状态
        if rtu_config.g_mqtt_connected:
            rtu_mqtt.publish_ota_status(data)
            
    except Exception as e:
        print("[OTA] Notify status error: {}".format(e))

def process_ota_command(command):
    """处理OTA命令"""
    try:
        action = command.get("action")
        
        if action == "check_update":
            result = check_update()
            return {
                "status": "success" if result else "failed",
                "update_available": result
            }
        elif action == "get_status":
            return {
                "status": "success",
                "ota_status": get_ota_status()
            }
        else:
            return {
                "status": "failed",
                "message": "Unknown OTA action"
            }
            
    except Exception as e:
        print("[OTA] Process command error: {}".format(e))
        return {
            "status": "failed",
            "message": str(e)
        }

def test_ota():
    """测试OTA功能"""
    print("[OTA] Testing OTA functionality...")
    
    # 测试初始化
    if init():
        print("[OTA] Init test passed")
    else:
        print("[OTA] Init test failed")
    
    # 测试检查更新
    print("[OTA] Check update test...")
    check_update()
    
    # 显示OTA状态
    status = get_ota_status()
    print("[OTA] OTA status: {}".format(status))
    
    print("[OTA] OTA test completed")